// WndSubclass.cpp: implementation of the CWndSubclass class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "WndSubclass.h"

#include <map>

using namespace std;

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif

//////////////////////////////////////////////////////////////////////////////////
// The message hook map is derived from CMapPtrToPtr, which associates			//	
// a pointer with another pointer. It maps an HWND to a CWndSubclass, like		//	
// the way MFC's internal maps map HWND's to CWnd's. The first CWndSubclass		//	
// attached to a window is stored in the map; all other CWndSubclass's for that	//	
// window are then chained via CWndSubclass::m_pNext.							//	
//
typedef std::map<ULONG, CWndSubclass*> WNDMAP;
class CWndSubclassMap  {
public:
	CWndSubclassMap();
	~CWndSubclassMap();
public:
	static CWndSubclassMap& GetHookMap();
	void Add(HWND hwnd, CWndSubclass* pWndSubclass);
	void Remove(CWndSubclass* pWndSubclass);
	void RemoveAll(HWND hwnd);
	CWndSubclass* Lookup(HWND hwnd);
private:
	WNDMAP m_WndMap;
};

// This trick is used so the hook map isn't
// instantiated until someone actually requests it.
// Please see the implement of GetHookMap()
#define	theHookMap	(CWndSubclassMap::GetHookMap())
//
//																				//																																								
//////////////////////////////////////////////////////////////////////////////////

/////////////////////////////////////////////////////////////////////////////////
// Subclassed window proc for message hooks. Replaces AfxWndProc (or whatever  //
// else was there before.)
//
LRESULT CALLBACK HookWndProc(HWND hwnd, UINT msg, WPARAM wp, LPARAM lp)
{
	// Set up MFC message state just in case anyone wants it
	// This is just like AfxCallWindowProc, but we can't use that because
	// a CWndSubclass is not a CWnd.
	//
// 	MSG& curMsg = AfxGetThreadState()->m_lastSentMsg;
// 	MSG  oldMsg = curMsg;   // save for nesting
// 	curMsg.hwnd		= hwnd;
// 	curMsg.message = msg;
// 	curMsg.wParam  = wp;
// 	curMsg.lParam  = lp;

	// Get hook object for this window. Get from hook map
	CWndSubclass* pWndSubclass = theHookMap.Lookup(hwnd);
	ASSERT(pWndSubclass);

	LRESULT lr;
	if (msg==WM_NCDESTROY) {
		// Window is being destroyed: unhook all hooks (for this window)
		// and pass msg to orginal window proc
		//
		WNDPROC wndproc = pWndSubclass->m_pOldWndProc;
		theHookMap.RemoveAll(hwnd);
		lr = ::CallWindowProc(wndproc, hwnd, msg, wp, lp);

	} else {
		// pass to msg hook
		lr = pWndSubclass->WindowProc(msg, wp, lp);
	}
	//curMsg = oldMsg;			// pop state
	return lr;
}
//																				//
//////////////////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////////////////
//Implementation of CWndSubclass														//

CWndSubclass::CWndSubclass()
{
	m_pOldWndProc = NULL;	
	m_pNext = NULL;
	m_hWndHooked  = NULL;
}

CWndSubclass::~CWndSubclass()
{
	ASSERT(m_hWndHooked==NULL);		// Can't destroy while still hooked!
	ASSERT(m_pOldWndProc==NULL);
}

/* Hook a window.
   This installs a new window proc that directs messages to the CWndSubclass.
   hWnd=NULL to remove.*/
BOOL CWndSubclass::SubclassWindow(HWND hWnd)
{
	// Hook the window
	ASSERT(m_hWndHooked==NULL);
	HWND hwnd = hWnd;

	ASSERT(hwnd && ::IsWindow(hwnd));
	theHookMap.Add(hwnd, this);			// Add to map of hooks
	m_hWndHooked = hWnd; 

	/*In order to repaint the new appearnce*/
	::SendMessage(m_hWndHooked, WM_NCPAINT, 1, 0);		//Change the interface at once
	//Invalidate()
	/*In order to repaint the new appearnce*/
	return TRUE;
}

/*Detach the hook*/
BOOL CWndSubclass::UnSubclassWindow()
{
	ASSERT(m_hWndHooked!=NULL);
	theHookMap.Remove(this);			// Remove from map
	m_pOldWndProc = NULL;

	/*In order to repaint the new appearnce*/
	::SendMessage(m_hWndHooked, WM_NCPAINT,1,0);		//Change the interface at once
	//Invalidate();
	/*In order to repaint the new appearnce*/

	m_hWndHooked = NULL; //pWnd is NULL
	return TRUE;
}

/* Window proc-like virtual function which specific CWndSubclasss will
   override to do stuff. Default passes the message to the next hook; 
   the last hook passes the message to the original window.
   You MUST call this at the end of your WindowProc if you want the real
   window to get the message. This is just like CWnd::WindowProc, except that
   a CWndSubclass is not a window.*/

LRESULT CWndSubclass::WindowProc(UINT msg, WPARAM wp, LPARAM lp)
{
	ASSERT(m_pOldWndProc);
	return m_pNext ? m_pNext->WindowProc(msg, wp, lp) :	
		::CallWindowProc(m_pOldWndProc, m_hWndHooked, msg, wp, lp);
}

/* Like calling base class WindowProc, but with no args, so individual
   message handlers can do the default thing. Like CWnd::Default*/

LRESULT CWndSubclass::Default()

{
	// MFC stores current MSG in thread state
	//MSG& curMsg = AfxGetThreadState()->m_lastSentMsg;
	// Note: must explicitly call CWndSubclass::WindowProc to avoid infinte
	// recursion on virtual function
	MSG curMsg;
	memset(&curMsg, 0, sizeof(MSG));
	return CWndSubclass::WindowProc(curMsg.message, curMsg.wParam, curMsg.lParam);
}

//																				//																																								
//////////////////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////////////////
//Implement of CWndSubclassMap													//
//
CWndSubclassMap::CWndSubclassMap()
{
}

CWndSubclassMap::~CWndSubclassMap()
{
	ASSERT(m_WndMap.size() == 0);	// All hooks should be removed!	
}

/* Get the one and only global hook map*/
CWndSubclassMap& CWndSubclassMap::GetHookMap()
{
	// By creating theMap here, C++ doesn't instantiate it until/unless
	// it's ever used! This is a good trick to use in C++, to
	// instantiate/initialize a static object the first time it's used.
	static CWndSubclassMap theMap;
	return theMap;
}

/* Add hook to map; i.e., associate hook with window*/
void CWndSubclassMap::Add(HWND hwnd, CWndSubclass* pWndSubclass)
{
	ASSERT(hwnd && ::IsWindow(hwnd));

	// Add to front of list
	pWndSubclass->m_pNext = Lookup(hwnd);
	m_WndMap[(ULONG)hwnd] = pWndSubclass;
	
	if (pWndSubclass->m_pNext==NULL) {
		// If this is the first hook added, subclass the window
		pWndSubclass->m_pOldWndProc = 
			(WNDPROC)SetWindowLong(hwnd, GWL_WNDPROC, (DWORD)HookWndProc);
	} else {
		// just copy wndproc from next hook
		pWndSubclass->m_pOldWndProc = pWndSubclass->m_pNext->m_pOldWndProc;
	}
	ASSERT(pWndSubclass->m_pOldWndProc);
}


/* Remove a hook class from map, but the wnd may still be subclassed*/
void CWndSubclassMap::Remove(CWndSubclass* pUnHook)
{
	HWND hwnd = pUnHook->m_hWndHooked;
	ASSERT(hwnd && ::IsWindow(hwnd));

	CWndSubclass* pHook = Lookup(hwnd);
	ASSERT(pHook);
	if (pHook==pUnHook) {
		// hook to remove is the first one in the hash table: replace w/next
		if (pHook->m_pNext)	{
			m_WndMap[(ULONG)hwnd] = pHook->m_pNext;
		} else {
			// This is the last hook for this window: restore wnd proc
			m_WndMap.erase((ULONG)hwnd);
			SetWindowLong(hwnd, GWL_WNDPROC, (DWORD)pHook->m_pOldWndProc);
		}
	} else {
		// Hook to remove is in the middle: just remove from linked list
		while (pHook->m_pNext!=pUnHook)
			pHook = pHook->m_pNext;
		ASSERT(pHook && pHook->m_pNext==pUnHook);
		pHook->m_pNext = pUnHook->m_pNext;
	}
}

/* Remove all the hooks for a window, the wnd is non-sublcassed */
void CWndSubclassMap::RemoveAll(HWND hwnd)
{
	CWndSubclass* pWndSubclass;
	while ((pWndSubclass = Lookup(hwnd))!=NULL){
		pWndSubclass->UnSubclassWindow();	// (unhook) this function will remove this point from MAP
	}
}

/* Find first hook associate with window*/
CWndSubclass* CWndSubclassMap::Lookup(HWND hwnd)
{
	CWndSubclass* pFound = NULL;
	
	WNDMAP::iterator it;
	it = m_WndMap.find((ULONG)hwnd);
	if(it != m_WndMap.end() )
	{
		pFound = (*it).second;
	}
	else
	{
		pFound = NULL;
	}
	return pFound;
}

//																				//																																								
//////////////////////////////////////////////////////////////////////////////////
